#include <iostream>
#include <cassert>
#include "raytracer.hh"
#include "surface_point.hh"
#include "sampling.hh"

using namespace std;
using namespace cgmath;

static void put_pixel (SDL_Surface *surface, int x, int y, const vector_3d& color);

raytracer::raytracer (void)
  : paths (8), branching_factor (8)
{}

raytracer::raytracer (double focal_distance_t)
  : use_dof (true),
    focal_distance (focal_distance_t),
    paths (8),
    branching_factor (8)
{}

object* raytracer::hit_object (const cgmath::ray_3d& ray, double max_l)
{
  object* obj = 0;
  for (unsigned i = 0; i < objects.size(); ++i)
    {
      double ti = objects[i]->hit (ray);
      if (ti > 1e-6 && ti < max_l)
	{
	  max_l = ti;
	  obj = objects[i];
	}
    }
  return obj;
}

object* raytracer::hit_surface (const cgmath::ray_3d& ray, 
				double max_l,
				surface_point& sp)
{
  object* obj = hit_object (ray, max_l);
  if (obj)
    obj->hit_surface_point(sp);
  return obj;
}

vector_3d raytracer::path_trace (const ray_3d& ray, unsigned depth)
{
  vector_3d color = vec (0.0, 0.0, 0.0);
  surface_point sp;
  if (! hit_surface (ray, numeric_limits<double>::infinity(), sp))
    return shade_exiting (ray);

  color += shade (normalized(-ray.d), sp);
  
  if (depth == 0)
    return color;

  bsdf_samples samples;
  sp.sample_bsdf (branching_factor, normalized(-ray.d), samples);


  vector_3d indirect_color = vec(0.0, 0.0, 0.0);
  for (unsigned i = 0; i < samples.size(); ++i)
    {

      ray_3d ray_sample (sp.point, samples[i].direction);
      indirect_color += mul(path_trace (ray_sample, depth - 1),
                            samples[i].weight);
    }
  indirect_color *= (1.0 / (double) branching_factor);
  color += indirect_color;
  
  return color;
}

cgmath::vector_3d raytracer::russian_path_trace (const cgmath::ray_3d& ray)
{
  vector_3d color = vec (0.0, 0.0, 0.0);
  surface_point sp;
  if (! hit_surface (ray, numeric_limits<double>::infinity(), sp))
    return shade_exiting (ray);

  color += shade (normalized(-ray.d), sp);

  // Determine if we should compute reflections.
  vector_3d indirect_color = vec (0.0, 0.0, 0.0);
  if ((double)rand() / (double)RAND_MAX > 0.5)
    {
      bsdf_samples samples;
      sp.sample_bsdf (1, normalized(-ray.d), samples);

      indirect_color = vec(0.0, 0.0, 0.0);
      for (unsigned i = 0; i < samples.size(); ++i)
	{
	  ray_3d ray_sample (sp.point, samples[i].direction);
	  indirect_color += mul(russian_path_trace (ray_sample),
				samples[i].weight);
	}
    }

  

  // The indirect color is multiplied by two because, the probability
  // of reflection is 0.5. Thus we gain an expectation value of 1 for
  // the computation of indirect lighting. (See exercise 11).
  return color + (2.0 * indirect_color);
}

vector_3d raytracer::trace (const ray_3d& ray)
{
  surface_point sp;
  if (hit_surface (ray, numeric_limits<double>::infinity(), sp))
    return shade (normalized(-ray.d), sp);
  else
    return shade_exiting (ray);
}

vector_3d raytracer::shade (const vector_3d& out, 
			    const surface_point& sp)
{
  vector_3d color (0.0);
  // Handle reflections from lights
  for (unsigned i = 0; i < lights.size(); ++i)
    {
      light_samples samples;
      lights[i]->illuminate (sp.point, 1, samples);

      // Here you should check that the light is not
      // occluded to get shadows.
      for (unsigned j = 0; j < samples.size(); ++j)
        {
          // Cast a ray from sp to light. todo
          ray_3d ray (sp.point, samples[j].direction);
          object* obj = hit_object(ray, samples[j].distance);

          if (obj == 0)
            {
              // No obstructions between point and light sample
              color += mul(sp.bsdf(-normalized(samples[j].direction), out), 
                           samples[j].radiance);
            }
        }
    }
  // In addition to direct lighting, the surface should
  // get light from other surfaces.
  return color;
}

vector_3d raytracer::shade_exiting (const cgmath::ray_3d& ray)
{
  // Color of ray that exits the scene.
  return vec(0.05, 0.05, 0.05);
}

void raytracer::trace_image (const matrix_4d& P, SDL_Surface* s)
{
  for (int y = 0; y < s->h; ++y)
    trace_scanline (P, s, y);
}

void raytracer::trace_scanline (const matrix_4d& P, 
                                SDL_Surface* s,
                                unsigned y)
{
  matrix_4d iP  = inv(P);
  vector_3d eye = cartesian (iP * vec(0.0, 0.0, 1.0, 0.0));
  for (int x = 0; x < s->w; ++x)
    {
      vector_3d p1 = cartesian(iP * vec<double>(x, y, 1.0, 1.0));
      put_pixel (s, x, y, trace (ray_3d (eye, p1 - eye)));
    }
}

// Note that P is from world to camera transformation.
void raytracer::dof_trace_scanline (const matrix_4d& P,
                                    const matrix_4d& PVP,
				    SDL_Surface* s,
				    unsigned y)
{
  matrix_4d iP  = inv(P);
  matrix_4d iPVP = inv(PVP); // to camera coordinate system
  //vector_3d eye = cartesian (iPVP * vec(0.0, 0.0, 1.0, 0.0));
  vector_3d eye = vec (0.0, 0.0, 0.0);

  //cout << eye << "\n";

  for (int x = 0; x < s->w; ++x)
    {

      // Trace a number of random rays through the pixel
      vector_3d color = vec (0.0, 0.0, 0.0);
      for (unsigned i = 0; i < paths; ++i)
        {
          double px = x;
          double py = y;
          
          vector_3d p1 = cartesian(iPVP * vec<double>(px, py, 1.0, 1.0));
          ray_3d r (eye, p1 - eye);

          //cout << eye << "\n";
          //cout << p1 << "\n";

          //focal_distance = 1000.0; // debug
          
          // Plane of focus
          vector_3d n_world = vec (0.0, 0.0, 1.0);
          vector_3d p_world = vec (0.0, 0.0, focal_distance);
          plane_3d plane_t (n_world, dot(p_world, n_world));
          
          // Intersection of the original ray with the plane of focus.
          double t;
          intersect (r, plane_t, t);
          vector_3d focus_point = r (t);

          //cout << t << "\n";
          //cout << focus_point << "\n";

          vector_3d jitter = 0.3 * lens_sample ();
          ray_3d lens_ray (eye + jitter, focus_point - (eye + jitter));

          // Transform the ray into world coordinate system
          lens_ray.p = cartesian (iP * vec (lens_ray.p(0),
                                            lens_ray.p(1),
                                            lens_ray.p(2),
                                            1.0));
          lens_ray.d = normal_transform(P, vec (lens_ray.d(0),
                                                lens_ray.d(1),
                                                lens_ray.d(2)));

          color += russian_path_trace (lens_ray);
        }
      
      color *= (1.0 / (double)paths);
      put_pixel(s, x, y, color);
    }
}

void raytracer::path_trace_scanline (const matrix_4d& P, 
                                     SDL_Surface* s,
                                     unsigned y,
				     bool russian)
{
  matrix_4d iP  = inv(P);
  vector_3d eye = cartesian (iP * vec(0.0, 0.0, 1.0, 0.0));
  for (int x = 0; x < s->w; ++x)
    {

      // Trace a number of random rays through the pixel
      vector_3d color = vec (0.0, 0.0, 0.0);
      for (unsigned i = 0; i < paths; ++i)
        {
          double dx = (double)rand() / (double)RAND_MAX;
          double dy = (double)rand() / (double)RAND_MAX;
          double px = x + dx;
          double py = y + dy;

	  vector_3d p1 = cartesian(iP * vec<double>(px, py, 1.0, 1.0));
	  ray_3d r (eye, p1 - eye);
	  
	  if (russian)
	    color += russian_path_trace (ray_3d (eye, p1 - eye));
	  else
	    color += path_trace (ray_3d (eye, p1 - eye), 1);

	}
      color *= (1.0 / (double)paths);

      put_pixel(s, x, y, color);
    }
}

void raytracer::clear_screen (SDL_Surface *s)
{
  vector_3d color = vec (0.0, 0.0, 0.0);
  for (int y = 0; y < s->h; ++y)
    for (int x = 0; x < s->w; ++x)
      put_pixel(s, x, y, color);
}

/// Sets pixel at x,y in surface to color. Adapted from SDL examples.
void put_pixel (SDL_Surface *surface, int x, int y, const vector_3d& color)
{
  // Gamma correct, clamp, and round color values to 8 bit.
  Uint8 r = round(255 * pow(max(0.0, min(color(0), 1.0)), 1.0/2.2));
  Uint8 g = round(255 * pow(max(0.0, min(color(1), 1.0)), 1.0/2.2));
  Uint8 b = round(255 * pow(max(0.0, min(color(2), 1.0)), 1.0/2.2));
  Uint32 pixel = SDL_MapRGBA (surface->format, r, g, b, 255);
  int bpp = surface->format->BytesPerPixel;
  Uint8 *p = (Uint8 *)surface->pixels + y * surface->pitch + x * bpp;
  switch(bpp) 
    {
    case 1:
      *p = pixel;
      break;
    case 2:
      *(Uint16 *)p = pixel;
      break;
    case 3:
      if(SDL_BYTEORDER == SDL_BIG_ENDIAN) {
	p[0] = (pixel >> 16) & 0xff;
	p[1] = (pixel >> 8) & 0xff;
	p[2] = pixel & 0xff;
      } else {
	p[0] = pixel & 0xff;
	p[1] = (pixel >> 8) & 0xff;
	p[2] = (pixel >> 16) & 0xff;
      }
      break;
    case 4:
      *(Uint32 *)p = pixel;
      break;
    }
}
